# !/usr/bin/env python
# -*-coding:utf-8 -*-
# Time       ：2021/12/6 10:07
# Author     ：caoxu
# version    ：python 3.8
# Description：测试yolov+ResNet特征提取、相似度计算
# from yolov3_detector.paddle_yolo import YOLO_v3
from common.resnet50 import CustomOperator
import cv2
import os
import shutil
import numpy as np
import scipy.spatial.distance as dist
# ignore warning log
import warnings
from yolov3_detector.pytorch_yolo.detect import detect_image
from common.config import OID_MODEL_PATH, OID_WEIGHT_PATH, OID_CLASS_PATH
from yolov3_detector.pytorch_yolo.models import load_model
from yolov3_detector.pytorch_yolo.utils.utils import load_classes
warnings.filterwarnings('ignore')
coco_yolo_detect = False
oid_yolo_detect = True
feature_extraction = False
similarity_calculate = False


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])


def euclidean(p, q):
    # 如果两数据集数目不同，计算两者之间都对应有的数
    # same = 0
    # for i in p:
    #     if i in q:
    #         same += 1
    # 计算欧几里德距离,并将其标准化
    e = sum([(p[i] - q[i]) ** 2 for i in range(len(p))])
    return 1 / (1 + e ** .5)


def yolo_recommend():
    if coco_yolo_detect:
        img_path = 'down_image/woman_bag.jpg'
        detect_obj_dir = 'detect_obj_dir/coco_yolo_detect/'
        if os.path.exists(detect_obj_dir):
            shutil.rmtree(detect_obj_dir)
        os.makedirs(detect_obj_dir, exist_ok=True)
        detector = YOLO_v3()
        # yolo目标检测生成子图片
        image_data = cv2.imread(img_path)
        detector.execute(image_data, detect_obj_dir)

    if oid_yolo_detect:
        img_path = 'down_image/Sweatshirt.jpg'
        detect_obj_dir = 'down_image/result/'
        if os.path.exists(detect_obj_dir):
            shutil.rmtree(detect_obj_dir)
        os.makedirs(detect_obj_dir, exist_ok=True)

        # load class
        detector_classes = load_classes(OID_CLASS_PATH)
        # load model
        detector_model = load_model(OID_MODEL_PATH, OID_WEIGHT_PATH)
        detector_model.eval()  # Set model to evaluation mode
        detect_image(model=detector_model,
                     image_path=img_path,
                     classes_names=detector_classes,
                     output_path=detect_obj_dir,
                     img_size=320,
                     conf_thres=0.1,
                     nms_thres=1)

    if feature_extraction:
        model = CustomOperator()
        # resnet50模型生成商品原图片向量
        vectors = []
        feature = model.execute(img_path)
        vectors.append(feature)

        # resnet50模型生成目标检测子图片向量
        obj_images = os.listdir(detect_obj_dir)
        obj_images.sort()
        for obj_image in obj_images:
            vector = model.execute(img_path + '/' + obj_image)
            vectors.append(vector)
        print('detect done')

    if similarity_calculate:
        query_img_path = 'down_image/women_cup.png'
        target_img_path = 'milvus_images/8911282'

        image_filenames = [os.path.join(target_img_path, x) for x in os.listdir(target_img_path) if is_image_file(x)]
        for image_filename in image_filenames:
            model = CustomOperator()
            # 特征抽取网络生成query image向量
            query_feature = model.execute(query_img_path)

            # 特征抽取网络生成target image向量
            target_feature = model.execute(image_filename)
            numver_list = []
            with open('milvus_images/8911282/8911282000_data.txt') as f:
                for line in f:
                    numver_list.extend([float(i) for i in line.split(',')])
            # 向量相似度计算
            # 欧几里德距离   衡量两个向量距离的远近
            # dis = dist.euclidean(query_feature, target_feature)
            # similarity_score = 1 / (1 + dis)

            # 内积距离   夹角余弦[-1,1]衡量两个向量方向的差异  similarity_score = cos(θ)
            dis = dist.cosine(query_feature, target_feature)
            similarity_score = 1 - dis
            dis2 = dist.cosine(query_feature, numver_list)
            similarity_score2 = 1 - dis2

            # 杰卡德距离

            # 汉明距离
            # dis = dist.hamming(query_feature, target_feature)
            # similarity_score = 1 - dis
            print(image_filename)
            print('distance', dis)
            print('score', similarity_score)
            print('distance2', dis2)
            print('score2', similarity_score2)


if __name__ == '__main__':
    yolo_recommend()
    print('done')
